package org.luosl.webmagicx.scheduler

import java.io.Closeable
import java.util.concurrent.locks.ReentrantLock
import java.util.concurrent.{LinkedBlockingQueue, PriorityBlockingQueue}

import org.apache.http.annotation.GuardedBy
import org.luosl.webmagicx.conf.{SpiderConf, XmlProps}
import us.codecraft.webmagic.{Request, Task}
import org.luosl.webmagicx.utils.ResourceUtils._

/**
  * 基于内存的优先级调度器
  * Created by luosl on 2018/3/15.
  */
class PriorityQueueScheduler(sc:SpiderConf, task:Task, props:XmlProps) extends AbstractScheduler(sc, task, props) with Closeable{


  private val lock:ReentrantLock = new ReentrantLock

  /**
    * 队列初始容量
    */
  private val initialCapacity:Int = 10
  /**
    * 优先级为0的队列
    */
  private val noPriorityQueue:LinkedBlockingQueue[Request] = new LinkedBlockingQueue[Request]()

  /**
    * 正优先级队列
    */
  private val priorityQueuePlus = new PriorityBlockingQueue[Request](initialCapacity, requestComparator)

  /**
    * 负优先级队列
    */
  private val priorityQueueMinus = new PriorityBlockingQueue[Request](initialCapacity, requestComparator)


  override def getLeftRequestsCount(task: Task): Int = noPriorityQueue.size()

  override def getTotalRequestsCount(task: Task): Int = getDuplicateRemover.getTotalRequestsCount(task)

  override def pushWhenNoDuplicate(request: Request, task: Task): Unit ={
    request.getPriority match {
      case 0 => noPriorityQueue.put(request)
      case num if num > 0 => priorityQueuePlus.put(request)
      case num if num < 0 => priorityQueueMinus.put(request)
    }
  }

  /**
    * get an url to crawl
    *
    * @param task the task of spider
    * @return the url to crawl
    */
  @GuardedBy("lock")
  override def poll(task: Task): Request = {
    withLock(lock){ _ =>
      val pr:Request = priorityQueuePlus.poll()
      if(null != pr) return pr
      val nr:Request = noPriorityQueue.poll()
      if(null != nr) return nr
      priorityQueueMinus.poll()
    }
  }

  override def close(): Unit = {
    noPriorityQueue.clear()
    priorityQueuePlus.clear()
    priorityQueueMinus.clear()
  }
}
